# -*- coding:utf-8 -*-
import os
import numpy as np

import torch
import torch.nn.parallel
import torch.optim

import nibabel as nib
from torch import nn
from torch.backends import cudnn

from DataPreProcess import preprocess


def main_test():

    ##############################
    # Using GPU
    ##############################
    # os.environ['CUDA_VISIBLE_DEVICES'] = '7'
    # model = torch.load('model_epoch_20.pkl')
    # model = nn.DataParallel(model)
    # cudnn.benchmark = True


    ##############################
    # Using CPU
    ##############################
    model = torch.load('model_epoch_20.pkl', map_location=lambda storage, loc: storage)
    model = model.module

    ##Load one patient
    dataPath = [os.path.abspath('.')+"/1875963"]

    ##Load patiens
    # dataPath = glob.glob(fileDir)

    savePath =os.path.abspath('.')
    data = preprocess(dataPath)
    print("Finish pre-processing")

    predProbMapsPerClass = []


    ####################################
    # When predict complete volumn
    # Top and bottom of complete label should use padding
    # They are considered no annotation
    ####################################

    # zeroPad = np.zeros((data[0].shape[1], data[0].shape[2]))
    # for i in range(2):
    #     predProbMapsPerClass.append(zeroPad)

    # for i in range(len(data)):
    for i in range(30, 35):  # For cpu test
        x = data[i]
        x = x[np.newaxis, :]
        output = test(x, model)
        output = output.data.numpy()
        predProbMapsPerClass.append(output)

    # for i in range(2):
    #     predProbMapsPerClass.append(zeroPad)

    predProbMapsPerClass = np.array(predProbMapsPerClass)
    predProbMapsPerClass = predProbMapsPerClass.reshape(predProbMapsPerClass.shape[0], predProbMapsPerClass.shape[2], predProbMapsPerClass.shape[3], predProbMapsPerClass.shape[4])

    predSegmentation = np.argmax(predProbMapsPerClass, axis=1)

    seg = np.zeros(predSegmentation.shape, dtype=np.int16)
    seg[predSegmentation == 1] = 1
    seg[predSegmentation == 2] = 2
    seg[predSegmentation == 3] = 3
    seg[predSegmentation == 4] = 4
    seg[predSegmentation == 5] = 5
    seg[predSegmentation == 6] = 6
    seg[predSegmentation == 7] = 7
    seg[predSegmentation == 8] = 8


    img = nib.Nifti1Image(seg, np.eye(4))
    nib.save(img, savePath + '/test.nii.gz')



def test(x,model):

    model.eval()

    x =torch.from_numpy(np.array(x))
    x1 = torch.autograd.Variable(x.float())

    ########GPU###########
    # x1 = torch.autograd.Variable(x.float()).cuda()
    ######################


    output = model(x1)  # nx5x9x9x9
    softmax = torch.nn.Softmax()
    output = output.view(-1, 9, 512 ** 2).permute(0, 2, 1).contiguous()
    output = output.view(-1, 9)


    output = softmax(output)
    output = output.view(-1, 512 ** 2, 9).permute(0, 2, 1).contiguous()
    output = output.view(-1, 9, 512, 512)
    return output


if __name__ == '__main__':
    main_test()

